import os
import torch
from PIL import Image
from torch.utils.data import Dataset

class CovidCTDataset(Dataset):

    def __init__(self, root_dir, txt_COVID, txt_NonCOVID, transform=None):

        self.root_dir = root_dir
        self.txt_path = [txt_COVID, txt_NonCOVID]
        self.classes = ['images/CT_COVID', 'images/CT_NonCOVID']  
        self.num_cls = len(self.classes)       
        #img_list二维数组每一张图片的路径和label
        self.img_list = []
        # 0是COVID，1是NonCOVID，在img_list中每一个数组的第二个元素保存着
        ## os.path.join 连接两个或更多的路径名
        for c in range(self.num_cls):
            cls_list = [[os.path.join(self.root_dir, self.classes[c], item), c] for item in read_txt(self.txt_path[c])]
            self.img_list += cls_list
        self.transform = transform

    def __len__(self):
        return len(self.img_list)

    # 图片的索引，0,1,2,3,4,,,
    def __getitem__(self, idx):
        if torch.is_tensor(idx):
            idx = idx.tolist()

        #得到第idx张图片的path
        img_path = self.img_list[idx][0]
        image = Image.open(img_path).convert('RGB')

        if self.transform:
            image = self.transform(image)
        # 
        sample = {'img': image,
                  'label': int(self.img_list[idx][1])}  # label里面就是0和1
        return sample


def read_txt(txt_path):
    with open(txt_path) as f:
        lines = f.readlines()
    txt_data = [line.strip() for line in lines]         # 主要是跳过'\n'
    return txt_data
